# 2048 Puzzle Game (Python Demo)

from typing import Tuple, List
import numpy as np
import random
import time


class Env:
    def __init__(self):
        self._init_trans_table()

    def _init_trans_table(self):
        """ init environment transition table"""
        move_left_table = np.zeros((17, 17, 17, 17, 4), dtype=np.uint16)
        move_right_table = np.zeros((17, 17, 17, 17, 4), dtype=np.uint16)
        score_table = np.zeros((17, 17, 17, 17), dtype=np.uint32)

        for a in range(17):
            for b in range(17):
                for c in range(17):
                    for d in range(17):
                        line = self._move_row([a, b, c, d])  # move left
                        move_left_table[a, b, c, d][::] = line
                        score_table[a, b, c, d] = self._score_line([a, b, c, d]) 

                        move_right_table[d, c, b, a][::] = line[::-1]  # move right

        self.move_left_table = move_left_table.tolist()
        self.move_right_table = move_right_table.tolist()
        self.score_table = score_table

    def _move_row(self, line):
        i = 0
        while(i<3):
            for j in range(i+1, 4):
                if line[j] != 0:
                    break
                j += 1
            if(j==4):
                break

            if(line[i]==0):
                line[i] = line[j]
                line[j] = 0
                if(j==3):
                    break
                else:
                    i -= 1
            elif(line[i]==line[j]):
                line[i] += 1
                line[j] = 0
            
            i += 1
        return line

    def _score_line(self, line):
        score = 0

        for tile in line:
            if tile>=2:
                score += (tile-1)*(2**tile)
        return score

    
    def print(self, board):
        state = '+' + '-' * 24 + '+\n'
        for row in [board[r:r + 4] for r in range(0, 16, 4)]:
            state += ('|' + ''.join('{0:6d}'.format((1 << t) & -2) for t in row) + '|\n')
        state += '+' + '-' * 24 + '+'
        print(board)
        print(state)

    def get_score(self, board, penalty):
        score_list = []
        for row in [board[r:r+4] for r in range(0, 16, 4)]:
            score_list.append(self.score_table[row[0], row[1], row[2], row[3]])
        return np.sum(score_list, dtype=np.uint32) - penalty
    
    def popup(self, _board:List[int])->Tuple[List[int], int]:
        _board = _board[:]
        empty = [i for i, t in enumerate(_board) if not t]

        tile = random.choice([1] * 9 + [2])
        _board[random.choice(empty)] = tile
        penalty_score = (tile>>1)*4
        return _board, penalty_score

    def reset(self):
        _board = [0] * 16
        _board, penalty_score = self.popup(_board)

        _board, p = self.popup(_board)
        penalty_score +=p
        return _board, penalty_score

    
    def mirror(self, _board:List[int])->List[int]:
        return [_board[r + i] for r in range(0, 16, 4) for i in reversed(range(4))]
    
    def transpose(self, _board:List[int])->List[int]:
        return [_board[r + i] for i in range(4) for r in range(0, 16, 4)]
    

    def left(self, _board:List[int])->List[int]:
        new_board = []
        for row in [_board[r:r+4] for r in range(0, 16, 4)]:
            new_row = self.move_left_table[row[0]][row[1]][row[2]][row[3]]
            new_board.extend(new_row)
        return new_board
   
    def right(self, _board:List[int])->List[int]:
        new_board = self.left(self.mirror(_board))
        return self.mirror(new_board)
    
    def up(self, _board:List[int])->List[int]:
        new_board = self.left(self.transpose(_board))
        return self.transpose(new_board)
    
    def down(self, _board:List[int])->List[int]:
        new_board = self.right(self.transpose(_board))
        return self.transpose(new_board)

    
if __name__ == '__main__':
    print('2048 Demo\n')

    env = Env()
    all_step = 0

    a=time.time()
    for game in range(10000):    
        board, penalty = env.reset()
        score = 0
        step = 0
        while True:
            # print('#{} [{}]'.format(step, score))
            # env.print(board)

            up_board = env.up(board)
            right_board = env.right(board)
            down_board = env.down(board)
            left_board = env.left(board)
            
            moves = [(up_board, env.get_score(up_board, penalty), "up"), 
                     (right_board, env.get_score(right_board, penalty), "right"), 
                     (down_board, env.get_score(down_board, penalty), "down"), 
                     (left_board, env.get_score(left_board, penalty), "left")]

            moves =[move for move in moves if move[0]!=board]

            # for label, move in zip(['up', 'right', 'down', 'left'], moves):
            #     print('{} = {}'.format(label, move[1]))
            # for move, _, label in moves:
            #     print('{} = {}'.format(label, move))
            
            if(not moves):
                all_step += step
                # print(game, all_step)
                break

            after, score, _ = max(moves, key = lambda move: move[1])
            # if reward == -1:
            #     break
            board, p = env.popup(after)
            # score += reward
            penalty += p
            step += 1
            # print()
    print(time.time()-a) 
